# -*- coding: utf-8 -*-

import os
import time
import h5py
import json
from milvus import Milvus, IndexType, MetricType

THRESHOLD = float(os.environ.get('THRESHOLD', '0.85'))  # 检索阈值
ID_MAP_FILE = "id_name_mapping.json"  # 本地映射文件

class MilvusRetrieval(object):
    def __init__(self, index_name, index_dir,
                 host=os.environ.get("MILVUS_HOST", "127.0.0.1"),
                 port=os.environ.get("MILVUS_PORT", 19530)):

        print(f"[INFO] Connecting to Milvus at {host}:{port} ...")
        t0 = time.time()
        self.client = Milvus(host, port)
        print(f"[INFO] ✅ Connected to Milvus ({time.time() - t0:.3f}s)\n")

        self.index_name = index_name
        self.index_dir = index_dir
        self.id_dict = {}
        self.load_or_create()

    def load_or_create(self):
        # 检查集合是否存在
        status, collections = self.client.list_collections()
        if self.index_name in collections:
            print(f"[INFO] 集合 '{self.index_name}' 已存在，复用已有集合...")
            self._load_existing_ids()
        else:
            print(f"[INFO] 集合 '{self.index_name}' 不存在，创建新集合并插入向量...")
            self._create_and_insert()

    def _load_existing_ids(self):
        """从本地 JSON 文件加载 ID 与 name 映射"""
        if os.path.exists(ID_MAP_FILE):
            with open(ID_MAP_FILE, 'r', encoding='utf-8') as f:
                self.id_dict = json.load(f)
            # JSON 中的 key 是 str，需要转回 int
            self.id_dict = {int(k): v for k, v in self.id_dict.items()}
            print(f"[INFO] 已加载本地 ID 映射，数量: {len(self.id_dict)}")
        else:
            print("[WARN] ID 映射文件不存在，检索结果将显示 'unknown' 名称")

    def _create_and_insert(self):
        # 1. 读取索引文件
        print(f"[INFO] 读取索引数据: {self.index_dir}")
        t0 = time.time()
        with h5py.File(self.index_dir, 'r') as h5f:
            self.retrieval_db = h5f['dataset_1'][:]
            self.retrieval_name = h5f['dataset_2'][:]
        print(f"[INFO] 读取完成，用时: {time.time() - t0:.3f}s, 向量数量: {len(self.retrieval_db)}, 维度: {self.retrieval_db.shape[1]}")

        # 2. 创建新集合
        t_create = time.time()
        self.client.create_collection({
            'collection_name': self.index_name,
            'dimension': self.retrieval_db.shape[1],
            'index_file_size': 1024,
            'metric_type': MetricType.IP
        })
        print(f"[INFO] 创建集合完成，用时: {time.time() - t_create:.3f}s")

        # 3. 插入向量
        print("[INFO] 插入向量到 Milvus ...")
        t_insert = time.time()
        status, ids = self.client.insert(
            collection_name=self.index_name,
            records=[v.tolist() for v in self.retrieval_db]
        )
        print(f"[INFO] 插入完成，用时: {time.time() - t_insert:.3f}s")

        # 4. 建立 ID 映射
        self.id_dict = {ids[i]: str(self.retrieval_name[i]) for i in range(len(ids))}

        # 5. 保存 ID 映射到本地
        with open(ID_MAP_FILE, 'w', encoding='utf-8') as f:
            json.dump(self.id_dict, f, ensure_ascii=False, indent=2)
        print(f"[INFO] 本地 ID 映射已保存: {ID_MAP_FILE}")

        # 6. 创建索引
        print("[INFO] 创建索引 (IndexType.FLAT) ...")
        t_index = time.time()
        self.client.create_index(self.index_name, IndexType.FLAT, {'nlist': 16384})
        print(f"[INFO] 索引创建完成，用时: {time.time() - t_index:.3f}s")
        print(f"[INFO] ✅ 索引加载完毕，总耗时: {time.time() - t0:.3f}s")

    def retrieve(self, query_vector, search_size=10):
        print(f"[INFO] 开始检索 top-{search_size} 向量...")
        t0 = time.time()

        status, vectors = self.client.search(
            collection_name=self.index_name,
            query_records=[query_vector],
            top_k=search_size,
            params={'nprobe': 16}
        )

        search_time = time.time() - t0
        print(f"[INFO] 检索完成，用时: {search_time:.3f}s")

        r_list = []
        for v in vectors[0]:
            score = float(v.distance) * 0.5 + 0.5
            if score > THRESHOLD:
                temp = {
                    "id": v.id,
                    "name": self.id_dict.get(v.id, "unknown"),
                    "score": round(score, 6)
                }
                r_list.append(temp)

        print(f"[INFO] 返回结果数量: {len(r_list)} (阈值: {THRESHOLD})\n")
        return r_list
